
from crewai.flow.flow import Flow, listen, start
from conn import mysql_conn,kb
from agents import get_news_agent,page_agent,get_entitys_agent,gen_entity_description

MC = mysql_conn.MysqlConn()

class OfflineFlow(Flow):

    @start()
    def get_news(self):
        res = get_news_agent.get_news()
        self.state['news'] =  res
        return res

    @listen(get_news)
    def get_page_content(self):
        news = self.state['news']
        for n in news:
            r = page_agent.get_page(n['url'])
            n['content_en'] = r['summary_en']
            n['content_zh'] = r['summary_zh']
            n['is_protected'] = r['is_protected']
            entitys = get_entitys_agent.entity_recognition(f'{n["title_zh"]}\n{n["content_zh"]}')
            for entity in entitys:
                if not entity['description']:
                    entity['description'] = gen_entity_description.gen_description(entity['name'])
            entity_ids = MC.insert('entitys',entitys)
            n['entitys']=','.join(entity_ids)
        self.state['news']=news
        return news


    @listen(get_page_content)
    def insert_into_mysql(self):
        print("insert into mysql")
        MC.insert('news',self.state['news'])
        print("mysql_done")
    #
    @listen(get_page_content)
    def insert_into_kb(self):
        print("insert into kb")
        kb.add_news_list(self.state['news'])
        print("kb_done")


if __name__ == "__main__":
    flow = OfflineFlow()
    flow.plot("plot/flow")
    a = flow.kickoff()
    print(a)